import numpy as np
import mindspore as ms
import mindspore.nn as nn
import mindspore.ops as ops
from mindspore.ops import functional as F
from mindspore.common.initializer import initializer


class TestNet(nn.Cell):
    def __init__(self):
        super(TestNet, self).__init__()

    def construct(self, input_x):
        shape1 = input_x.shape[1:]
        test_net = nn.LayerNorm(shape1, begin_norm_axis=1, begin_params_axis=1)
        out = test_net(input_x)
        return [out]


if __name__ == "__main__":
    net = TestNet()

    input_x = ms.Tensor(np.ones([20, 5, 10, 10]), ms.float32)
    result = net(input_x)
    print(result[0])
    print(result[0].shape)
